""" Package, 一些通用的计算函数 """
from __future__ import annotations
from pprint import pprint
import numpy as np
import pandas as pd

def cal_hourly(data: pd.DataFrame | pd.Series, time_resolution: str, std=False) -> tuple:
    """ # 小时均值计算,
        
        根据环境空气质量评价技术规范HJ 663——2013, https://www.mee.gov.cn/ywgz/fgbz/bz/bzwb/jcffbz/201309/t20130925_260809.shtml
        小时均值的统计是向后平均, 即01:00的数据是00:00-01:00的数据的均值
        
        Pandas计算均值时则是向前平均, 即01:00的数据是01:00-02:00的数据的均值
        因此, 使用Pandas计算完成后还需要补偿1小时, 即将时间索引DatetimeIndex + 1小时

        data: pd.DataFrame | pd.Series, 含有datetime时间索引, 分钟/小时分辨率;
        time_resolution: data的时间分辨率(很重要), 'T'代表分钟, 'S'代表秒, 实际数据是5秒、10秒的分辨率选择'S', 以此类推;
        std: 是否计算标准偏差, 默认: False

        return: (mean, std) -> (pd.DataFrame | pd.Series, pd.DataFrame | pd.Series | None)
    
    单进程
    2023-08-30 v1
    2023-10-07 v1.1 修复类型判断操作符"|"仅在python≥3.10下可用的bug
    """

    # 判断时间分辨率
    if time_resolution not in ["T", "S"]:
        raise ValueError("time_resolution错误: %s, 可选: T | S" % time_resolution)

    # 判断data类型
    # if not isinstance(data, pd.DataFrame | pd.Series):
    if not isinstance(data, (pd.DataFrame, pd.Series)):
        raise ValueError("data类型错误: %s, 可选: pd.DataFrame | pd.Series" % type(data))

    """ 计算小时均值 """
    if time_resolution == 'S':
        df_1min = data.resample("T").mean()
    elif time_resolution == 'T':
        df_1min = data
    else:
        raise ValueError()

    # 小时均值group
    group_1h = df_1min.resample("H")

    # 计算均值
    df_mean = group_1h.mean()
    
    # 时间补偿
    df_mean.index += pd.Timedelta(value=1, unit="H")

    # 计算标准偏差
    if std is not None:
        df_std = group_1h.std()
        df_std.index += pd.Timedelta(value=1, unit="H")
    else:
        df_std = None
    
    # 返回数据
    return df_mean, df_std


def cal_mda8(data: pd.DataFrame | pd.Series, time_resolution: str, min_num_moving_avg=6, min_num_8h=14) -> pd.DataFrame | pd.Series:
    """ # 日最大8小时平均, Mean values of the daily maximum 8 h average (MDA8)

        根据环境空气质量评价技术规范HJ 663——2013, https://www.mee.gov.cn/ywgz/fgbz/bz/bzwb/jcffbz/201309/t20130925_260809.shtml
        小时均值的统计是向后平均, 即01:00的数据是00:00-01:00的数据的均值;
        日最大8小时滑动平均计算方式如下:
            计算一天中08:00-24:00之间17个8小时滑动平均(C8, ..., C24), 08:00时刻的8小时滑动平均为01:00-08:00的8小时均值(C8)，
            取这17个均值中的最大值作为MDA8, 质控:17个中需要有至少14个有效数据, 每个时刻的8小时滑动平均至少需要6个有效数据;

            点位8小时滑动平均:
                N1 = average of {C18(前一天), ... , C23(前一天), C24(前一天), C1}
                ...
                N8 = average of {C1, C2, ..., C8}, N8表示08:00时的8小时滑动平均, C1表示01:00时的浓度
                ...
                N24 = average of {C17, C18, ..., C24}
            点位日最大8小时平均:
                MDA8 = max of {N8, ..., N24}

        统计有效性标准参考环境空气质量标准 GB 3095——2012, https://www.mee.gov.cn/ywgz/fgbz/bz/bzwb/dqhjbh/dqhjzlbz/201203/t20120302_224165.shtml

        data: pd.DataFrame, 含有datetime时间索引, 秒/分钟/小时分辨率;
        time_resolution: data的时间分辨率(很重要), 'H'代表小时, 'T'代表分钟, 'S'代表秒;
        min_num_moving_avg: 计算8h滑动平均时至少含有的数据个数;
        min_num_8h: 统计MDA8时至少需要包含的8h滑动均值数量;

        return: pd.DataFrame | pd.Series

    单进程
    2023-08-30 v1
    2023-10-07 v1.1 修复类型判断操作符"|"仅在python≥3.10下可用的bug
    """

    # 判断时间分辨率
    if time_resolution not in ["H", "T", "S"]:
        raise ValueError("time_resolution错误: %s, 可选: H | T | S" % time_resolution )

    # 判断data类型
    # if not isinstance(data, pd.DataFrame | pd.Series):
    if not isinstance(data, (pd.DataFrame, pd.Series)):
        raise ValueError("data类型错误: %s, 可选: pd.DataFrame | pd.Series" % type(data))

    """ 计算小时均值 """
    if time_resolution in ["T", 'S']:
        data, _ = cal_hourly(data=data, time_resolution=time_resolution, std=False)

    """ Pandas中DatetimeIndex索引默认一天24h为00:00-23:00, 0:00的浓度应该是前一天23:00-0:00之间的均值, 也即前一天的值
        此处需要先将时间索引 - 1小时, 即00:00 - 23:00的值实际表示的是1:00 - 24:00的小时值 """

    # 时间补偿
    data.index -= pd.Timedelta(value=1, unit="H")

    # 计算8小时滑动平均, 需要至少6个有效值
    df_8h_moving_average = data.rolling(window=8, min_periods=min_num_moving_avg).mean()

    # 删除00:00-06:00的所有数据, 即原本01:00-07:00的8小时滑动平均数据
    df_8h_moving_average[df_8h_moving_average.index.hour.isin(range(7))] = None

    """ 计算日最大8小时平均, 每日需满足14个滑动平均值 """
    # 按日分组
    group_8h = df_8h_moving_average.resample("D")

    # 统计每列非NaN值的个数
    df_count_no_nan = group_8h.count()

    # 计算日最大8小时均值
    df_mda8 = group_8h.max()

    # 筛选其中min_num_8h符合条件的数据
    df_mda8 = df_mda8[df_count_no_nan >= min_num_8h]

    # 删除空行
    df_mda8.dropna(axis=0, how="all", inplace=True)

    # 返回数据
    return df_mda8


# def average_daily(data: pd.DataFrame, time_resolution: str, min_num_daily=20):
    """ # 日均值
        参考 环境空气质量评价技术规范HJ 663——2013, https://www.mee.gov.cn/ywgz/fgbz/bz/bzwb/jcffbz/201309/t20130925_260809.shtml
        环境空气质量标准 GB 3095——2012, https://www.mee.gov.cn/ywgz/fgbz/bz/bzwb/dqhjbh/dqhjzlbz/201203/t20120302_224165.shtml
        每日有效数据包含至少20个小时值;

    data: pd.DataFrame, 含有datetime时间索引, 分钟/小时分辨率;
    time_resolution: data的时间分辨率(很重要), 'D'代表天, 'H'代表小时, 'T'代表分钟, 'S'代表秒;
    min_num_daily: 计算日均值时至少需要包含的小时值数量, 默认为20;

    2022-11-11    v2
    """

    # 判断时间分辨率
    if time_resolution not in ['D', 'H', 'T', 'S']:
        print('average_daily, time_resolution错误!')
        return None

    # 判断data类型
    if not isinstance(data, pd.DataFrame):
        print('average_daily, data不是DataFrame!')
        return None

    if time_resolution == 'D':
        return data.resample('D').mean()

    ''' 计算小时均值 '''
    df_1h = cal_hourly(data=data, time_resolution=time_resolution)

    ''' Pandas中DatetimeIndex索引默认一天24h为00:00-23:00, 
        但根据规定, 污染物浓度的小时均值数据是表示前一个小时的浓度平均;
        计算日均值浓度时0:00的浓度应该是前一天23:00-0:00之间的均值，
        也就是应该算到前一天的日均值中去, 这里将时间索引减去一个小时,
        转换为当前时刻代表后一个小时的均值,
        而当输入数据的分辨率是分钟或秒时则不存在这个问题;
        '''
    df_1h.index = df_1h.index - pd.Timedelta(value=1, unit='H')

    ''' 计算日平均, 每日需至少20个小时值 '''
    # 按日分组
    group_1day = df_1h.resample('D')

    # 查找小于20个小时值的日期索引
    dt_num_less_than_20 = [i[0] for i in group_1day if np.sum(~np.isnan(i[1].iloc[:, 0].to_numpy())) < min_num_daily]

    # 计算日均值
    df_1day = group_1day.mean()

    # 删除日有效数据小于20的天
    df_1day.loc[dt_num_less_than_20, :] = None

    return df_1day


def cal_daily(data: pd.DataFrame | pd.Series, time_resolution: str, min_num_daily=20, std=False):
    """ # 日均值计算,
        参考 环境空气质量评价技术规范HJ 663——2013, https://www.mee.gov.cn/ywgz/fgbz/bz/bzwb/jcffbz/201309/t20130925_260809.shtml
        环境空气质量标准 GB 3095——2012, https://www.mee.gov.cn/ywgz/fgbz/bz/bzwb/dqhjbh/dqhjzlbz/201203/t20120302_224165.shtml
        每日有效数据包含至少20个小时值;

        data: pd.DataFrame, 含有datetime时间索引, 秒/分钟/小时分辨率;
        time_resolution: data的时间分辨率(很重要), 'H'代表小时, 'T'代表分钟, 'S'代表秒;
        min_num_daily: 计算日均值时至少需要包含的小时值数量, 默认为20;
        std: 是否统计标准偏差

        return (mean, std) -> (pd.DataFrame | pd.Series, pd.DataFrame | pd.Series | None)

    单进程
    2023-08-30 v1
    2023-10-07 v1.1 修复类型判断操作符"|"仅在python≥3.10下可用的bug
    """
    # 判断时间分辨率
    if time_resolution not in ["H", "T", "S"]:
        raise ValueError("time_resolution错误: %s, 可选: H | T | S" % time_resolution )

    # 判断data类型
    # if not isinstance(data, pd.DataFrame | pd.Series):
    if not isinstance(data, (pd.DataFrame, pd.Series)):
        raise ValueError("data类型错误: %s, 可选: pd.DataFrame | pd.Series" % type(data))

    """ 计算小时均值 """
    if time_resolution != 'H':
        data, _ = cal_hourly(data=data, time_resolution=time_resolution, std=std)

    """ Pandas中DatetimeIndex索引默认一天24h为00:00-23:00, 0:00的浓度应该是前一天23:00-0:00之间的均值, 也即前一天的值
        此处需要先将时间索引 - 1小时, 即00:00 - 23:00的值实际表示的是1:00 - 24:00的小时值 """

    # 时间补偿
    data.index -= pd.Timedelta(value=1, unit="H")

    # 按日分组
    group_1day = data.resample("D")

    """ 计算日均值 """
    # 日均值包含的小时值数量
    df_count_no_nan = group_1day.count()

    # 日均值
    df_mean = group_1day.mean()

    # 筛选min_num_daily符合条件的索引
    df_bool = df_count_no_nan >= min_num_daily

    # 筛选
    df_mean = df_mean[df_bool]

    # 删除空行
    df_mean.dropna(axis=0, how="all", inplace=True)

    """ 计算标准偏差 """    
    if std:
        df_std = group_1day.std()
        df_std = df_std[df_bool]
        df_std.dropna(axis=0, how="all", inplace=True)
    else:
        df_std = None

    # 返回数据
    return df_mean, df_std


def cal_monthly(data: pd.DataFrame | pd.Series, time_resolution: str, min_num_monthly=27, percentile=False, std=False, **kwargs):
    """ # 月均值计算,
        
        参考 环境空气质量评价技术规范HJ 663——2013, https://www.mee.gov.cn/ywgz/fgbz/bz/bzwb/jcffbz/201309/t20130925_260809.shtml;
        环境空气质量标准 GB 3095——2012, https://www.mee.gov.cn/ywgz/fgbz/bz/bzwb/dqhjbh/dqhjzlbz/201203/t20120302_224165.shtml;

        每月的日均值有效数据至少为27天(2月为25天);

        data: pd.DataFrame, 含有datetime时间索引, 秒/分钟/小时/天分辨率;
        time_resolution: data的时间分辨率(很重要), 'D'代表天, 'H'代表小时, 'T'代表分钟, 'S'代表秒;
        min_num_monthly: 计算月均值时至少需要包含的日均值数量, 默认为27(2月为25);
        min_num_daily: 计算日均值时至少需要包含的小时值数量, 默认为20;
        percentile: 统计均值时的百分位数, 0-100, 0表示最小值, 100表示最大值, 90表示90分位数, 默认为False(即求均值);
        std: 不统计百分位数时是否统计标准偏差

        return (mean, std, percentile) -> (pd.DataFrame | pd.Series | None, 
                                            pd.DataFrame | pd.Series | None, 
                                            pd.DataFrame | pd.Series | None)

    单进程
    2023-08-30 v1
    2023-10-07 v1.1 修复类型判断操作符"|"仅在python≥3.10下可用的bug
    """

    # 判断时间分辨率
    if time_resolution not in ["D", "H", "T", "S"]:
        raise ValueError("time_resolution错误: %s, 可选: D | H | T | S" % time_resolution )

    # 判断data类型
    # if not isinstance(data, pd.DataFrame | pd.Series):
    if not isinstance(data, (pd.DataFrame, pd.Series)):
        raise ValueError("data类型错误: %s, 可选: pd.DataFrame | pd.Series" % type(data))

    # 计算日均值
    if time_resolution != 'D':
        if "min_num_daily" in kwargs.keys():
            data, _ = cal_daily(data=data,
                                    time_resolution=time_resolution,
                                    min_num_daily=kwargs["min_num_daily"],
                                )
        else:
            raise ValueError("min_num_daily参数不存在!")
    
    """ 月统计 """
    # 按月分组
    group_1month = data.resample("MS")

    # 统计每月非Nan值的个数
    df_count_no_nan = group_1month.count()

    # 筛选min_num_monthly符合条件的数据索引
    df_bool = df_count_no_nan >= min_num_monthly
    df_bool_feb = df_count_no_nan >= min_num_monthly - 2
    df_bool[df_bool.index.month == 2] = df_bool_feb[df_bool_feb.index.month == 2]

    # 统计月均值
    if percentile:

        # 分位数
        df_percentile = group_1month.quantile(q=percentile / 100)

        # 筛选min_num_monthly符合条件的数据
        df_percentile = df_percentile[df_bool]

        # 返回数据
        return None, None, df_percentile

    else:
        # 计算均值
        df_mean = group_1month.mean()

        # 筛选min_num_monthly符合条件的数据
        df_mean = df_mean[df_bool]

        # 计算标准偏差
        if std:
            df_std = group_1month.std()
            df_std = df_std[df_bool]
        else:
            df_std = None
        
        # 返回数据
        return df_mean, df_std, None


def cal_seasonal(data: pd.DataFrame | pd.Series, time_resolution: str, min_num_seasonal=60, percentile=False, std=False, **kwargs):
    """ # 季节均值计算,
        没有参考文献, 设定每季的日均值有效数据至少为60个;

        data: pd.DataFrame, 含有datetime时间索引,  秒/分钟/小时/天分辨率;
        time_resolution: data的时间分辨率(很重要), 'D'代表天, 'H'代表小时, 'T'代表分钟, 'S'代表秒;
        min_num_seasonal: 计算季节均值时至少需要包含的日均值数量, 默认为60;
        min_num_daily: 计算日均值时至少需要包含的小时值数量, 默认为20;
        percentile: 统计均值时的百分位数, 0-100, 0表示最小值, 100表示最大值, 90表示90分位数, 默认为False(即求均值);
        std: 不统计百分位数时是否统计标准偏差

        return (mean, std, percentile) -> (pd.DataFrame | pd.Series | None, 
                                            pd.DataFrame | pd.Series | None, 
                                            pd.DataFrame | pd.Series | None)

    单进程
    2023-08-30 v1
    2023-10-07 v1.1 修复类型判断操作符"|"仅在python≥3.10下可用的bug
    """

    # 判断时间分辨率
    if time_resolution not in ["D", "H", "T", "S"]:
        raise ValueError("time_resolution错误: %s, 可选: D | H | T | S" % time_resolution )

    # 判断data类型
    # if not isinstance(data, pd.DataFrame | pd.Series):
    if not isinstance(data, (pd.DataFrame, pd.Series)):
        raise ValueError("data类型错误: %s, 可选: pd.DataFrame | pd.Series" % type(data))

    # 计算日均值
    if time_resolution != 'D':
        if "min_num_daily" in kwargs.keys():
            data, _ = cal_daily(data=data,
                                    time_resolution=time_resolution,
                                    min_num_daily=kwargs["min_num_daily"],
                                )
        else:
            raise ValueError("min_num_daily参数不存在!")
    
    """ 季节统计 """
    # 按季节分组
    group_1season = data.resample("QS-MAR")

    # 查找日均值数据小于60个的季节
    df_bool = group_1season.count() >= min_num_seasonal

    # 统计季节均值
    if percentile:
        # 分位数
        df_percentile = group_1season.quantile(q=percentile / 100)

        # 删除日均值数据小于60个的季节
        df_percentile = df_percentile[df_bool]

        return None, None, df_percentile

    else:
        # 计算均值
        df_mean = group_1season.mean()

        # 筛选日均值数据大于60个的季节
        df_mean = df_mean[df_bool]

        # 计算标准偏差
        if std:
            df_std = group_1season.std()
            df_std = df_std[df_bool]
        else:
            df_std = None
        
        # 返回数据
        return df_mean, df_std, None


def cal_yearly(data: pd.DataFrame | pd.Series, time_resolution: str, min_num_annual=324, percentile=False, std=False,**kwargs):
    """ # 年均值计算
        参考 环境空气质量评价技术规范HJ 663——2013, https://www.mee.gov.cn/ywgz/fgbz/bz/bzwb/jcffbz/201309/t20130925_260809.shtml
        环境空气质量标准 GB 3095——2012, https://www.mee.gov.cn/ywgz/fgbz/bz/bzwb/dqhjbh/dqhjzlbz/201203/t20120302_224165.shtml

        每年的日均值有效数据至少为324个;

        data: pd.DataFrame, 含有datetime时间索引, 秒/分钟/小时/天分辨率;
        time_resolution: data的时间分辨率(很重要), 'D'代表天, 'H'代表小时, 'T'代表分钟, 'S'代表秒;
        min_num_annual: 计算年节均值时至少需要包含的日均值数量, 默认为324;
        min_num_daily: 计算日均值时至少需要包含的小时值数量, 默认为20;
        percentile: 统计均值时的百分位数, 0-100, 0表示最小值, 100表示最大值, 90表示90分位数, 默认为False(即求均值);
        std: 不统计百分位数时是否统计标准偏差

        return (mean, std, percentile) -> (pd.DataFrame | pd.Series | None, 
                                            pd.DataFrame | pd.Series | None, 
                                            pd.DataFrame | pd.Series | None)
    单进程
    2023-08-30 v1
    2023-10-07 v1.1 修复类型判断操作符"|"仅在python≥3.10下可用的bug
    """

    # 判断时间分辨率
    if time_resolution not in ["D", "H", "T", "S"]:
        raise ValueError("time_resolution错误: %s, 可选: D | H | T | S" % time_resolution )

    # 判断data类型
    # if not isinstance(data, pd.DataFrame | pd.Series):
    if not isinstance(data, (pd.DataFrame, pd.Series)):
        raise ValueError("data类型错误: %s, 可选: pd.DataFrame | pd.Series" % type(data))

    # 计算日均值
    if time_resolution != 'D':
        if "min_num_daily" in kwargs.keys():
            data, _ = cal_daily(data=data,
                                    time_resolution=time_resolution,
                                    min_num_daily=kwargs["min_num_daily"],
                                )
        else:
            raise ValueError("min_num_daily参数不存在!")

    ''' 计算年平均 '''
    # 按年分组
    # print('data:', data)
    group_1year = data.resample('YS')

    # 查找日均值数据小于324个的年份
    dt_num_less_than_324 = [i[0] for i in group_1year if np.sum(~np.isnan(i[1].iloc[:, 0].to_numpy())) < min_num_annual]

    # 统计百分位数
    if percentile:

        # 分位数
        df_percentile = group_1year.quantile(q=percentile / 100)

        # 筛选
        df_percentile.loc[dt_num_less_than_324, :] = None

        # 去除nan行
        # df_percentile.dropna(how='all', axis=0, inplace=True)

        # 返回数据
        return None, None, df_percentile
    
    else:
        # 计算均值
        df_mean = group_1year.mean()
        # print('df_mean:', df_mean)
        
        # 筛选
        df_mean.loc[dt_num_less_than_324, :] = None

        # 去除nan行
        # df_mean.dropna(how='all', axis=0, inplace=True)

        # 计算标准偏差
        if std:
            df_std = group_1year.std()
            df_std.loc[dt_num_less_than_324, :] = None
            
            # 去除nan行
            # df_std.dropna(how='all', axis=0, inplace=True)
            
        else:
            df_std = None
        
        # 返回数据
        return df_mean, df_std, None


def average_range_day(data: pd.DataFrame | pd.Series, percentile=False):
    """ # 计算输入数据的均值

    data: 时间分辨率必须为天
    percentile: 统计均值时的百分位数, 0-100, 0表示最小值, 100表示最大值, 90表示90分位数, 默认为False(即求均值);

    2022-11-12 v1
    单进程
    """

    # 统计分位数
    if percentile:
        # 分位数
        df_percentile = data.quantile(q=percentile / 100)

        return {"percentile": df_percentile}

    else:
        # 求平均
        df_mean = data.mean()
        df_std = data.std()

        return {"mean": df_mean, "std": df_std}


if __name__ == '__main__':

    # path_file = 'D:/Desktop/df_all.csv'
    # df = pd.read_csv(path_file, encoding='utf-8')
    # df['datetime'] = pd.to_datetime(df['datetime'])
    # df.set_index('datetime', inplace=True)
    # # print(df, '\n', df.head(48))
    # # average_daily(data=df)
    # average_monthly(data=df, time_resolution='H', min_num_daily=27, min_num_hourly=20)
    # # average_monthly(data=df, time_resolution='H', min_num_daily=27, args=(1, 3))


    df = pd.Series(data=range(100), index=pd.date_range('20220101', periods=100, freq='H'), name='a')
    print(df)
    a = cal_mda8(data=df, time_resolution='T', min_num_8h=14, min_num_moving_avg=6)
    pprint(a)
